"""轨道递推模块 - 支持不同的轨道递推方法"""

import numpy as np
from astropy import units as u
from astropy.time import Time, TimeDelta
from poliastro.bodies import Earth
from poliastro.twobody import Orbit
from poliastro.twobody.propagation import (
    FarnocchiaPropagator,
    ValladoPropagator,
    CowellPropagator
)
from poliastro.core.propagation import func_twobody
from poliastro.core.perturbations import J2_perturbation
from poliastro.earth import EarthSatellite
from poliastro.earth.enums import EarthGravity
from enum import Enum, auto

class PropagationMethod(Enum):
    """轨道递推方法"""
    KEPLER = auto()       # 二体模型，开普勒轨道
    J2 = auto()           # 考虑J2摄动
    HPOP = auto()         # 高精度轨道递推

class OrbitPropagator:
    """轨道递推器类"""
    
    def __init__(self, orbit, method=PropagationMethod.KEPLER):
        """
        初始化轨道递推器
        
        参数:
            orbit: poliastro轨道对象
            method: 递推方法
        """
        self.orbit = orbit
        self.method = method
        
    def propagate(self, time_of_flight):
        """
        递推轨道到指定时间
        
        参数:
            time_of_flight: 递推时间
            
        返回:
            新的轨道对象
        """
        if self.method == PropagationMethod.KEPLER:
            # 二体模型，使用Farnocchia传播器
            return self.orbit.propagate(time_of_flight, FarnocchiaPropagator())
            
        elif self.method == PropagationMethod.J2:
            # 考虑J2摄动的传播
            def f(t0, state, k):
                du_kep = func_twobody(t0, state, k)
                ax, ay, az = J2_perturbation(t0, state, k, J2=Earth.J2.value, R=Earth.R.to(u.km).value)
                du_ad = np.array([0, 0, 0, ax, ay, az])
                return du_kep + du_ad
                
            return self.orbit.propagate(time_of_flight, CowellPropagator(f=f))
            
        elif self.method == PropagationMethod.HPOP:
            # 高精度轨道递推 - 创建地球卫星对象并考虑多种摄动
            from poliastro.spacecraft import Spacecraft
            earth_satellite = EarthSatellite(self.orbit, Spacecraft(1.0 * u.kg, 1.0 * u.m ** 2))
            
            # 考虑地球引力场、大气阻力等多种摄动
            return earth_satellite.propagate(
                time_of_flight,
                gravity=EarthGravity.J2
            ).orbit
            
        else:
            raise ValueError("不支持的传播方法")

    def propagate_multiple(self, times):
        """
        递推轨道到多个时间点
        
        参数:
            times: 时间数组
            
        返回:
            轨道对象数组
        """
        if isinstance(times, u.Quantity):
            # 如果传入的是时间数量对象
            return [self.propagate(time) for time in times]
        else:
            # 如果传入的是TimeDelta或其他时间对象
            return [self.propagate(time - self.orbit.epoch) for time in times] 